"""
数据加载模块
负责从各种数据源加载意图识别数据
"""

import pandas as pd
import numpy as np
from typing import Tuple, Dict, List
from pathlib import Path
import logging
from sklearn.model_selection import train_test_split

logger = logging.getLogger(__name__)


class IntentDataLoader:
    """意图识别数据加载器"""
    
    def __init__(self, data_path: str = None):
        """
        初始化数据加载器
        
        Args:
            data_path: 数据文件路径
        """
        self.data_path = data_path
        self.df = None
        
    def load_from_csv(self, file_path: str) -> pd.DataFrame:
        """
        从CSV文件加载数据
        
        Args:
            file_path: CSV文件路径
            
        Returns:
            pd.DataFrame: 加载的数据
        """
        try:
            logger.info(f"正在从CSV文件加载数据: {file_path}")
            self.df = pd.read_csv(file_path)
            logger.info(f"数据加载成功，共 {len(self.df)} 条记录")
            return self.df
        except FileNotFoundError:
            logger.error(f"文件未找到: {file_path}")
            raise
        except pd.errors.EmptyDataError:
            logger.error(f"CSV文件为空: {file_path}")
            raise
        except Exception as e:
            logger.error(f"加载数据时发生错误: {e}")
            raise
    
    def validate_data(self, df: pd.DataFrame) -> bool:
        """
        验证数据格式
        
        Args:
            df: 数据框
            
        Returns:
            bool: 验证是否通过
        """
        required_columns = ['text', 'intent']
        
        # 检查必需列
        missing_columns = [col for col in required_columns if col not in df.columns]
        if missing_columns:
            logger.error(f"缺少必需列: {missing_columns}")
            return False
        
        # 检查空值
        null_counts = df[required_columns].isnull().sum()
        if null_counts.any():
            logger.warning(f"发现空值: {null_counts.to_dict()}")
        
        # 检查数据类型
        if not df['text'].dtype == 'object':
            logger.warning("text列不是字符串类型")
        
        if not df['intent'].dtype == 'object':
            logger.warning("intent列不是字符串类型")
        
        logger.info("数据验证完成")
        return True
    
    def get_data_statistics(self, df: pd.DataFrame) -> Dict:
        """
        获取数据统计信息
        
        Args:
            df: 数据框
            
        Returns:
            Dict: 统计信息
        """
        stats = {
            'total_samples': len(df),
            'num_classes': df['intent'].nunique(),
            'class_distribution': df['intent'].value_counts().to_dict(),
            'avg_text_length': df['text'].str.len().mean(),
            'min_text_length': df['text'].str.len().min(),
            'max_text_length': df['text'].str.len().max(),
            'null_counts': df.isnull().sum().to_dict(),
            'duplicate_texts': df['text'].duplicated().sum()
        }
        
        logger.info(f"数据统计: {stats}")
        return stats
    
    def split_data(self, df: pd.DataFrame, 
                  test_size: float = 0.3, 
                  val_size: float = 0.5,
                  random_state: int = 42) -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
        """
        分割数据集
        
        Args:
            df: 原始数据
            test_size: 测试集比例
            val_size: 在剩余数据中验证集的比例  
            random_state: 随机种子
            
        Returns:
            Tuple: (训练集, 验证集, 测试集)
        """
        logger.info(f"开始分割数据集，测试集比例: {test_size}, 验证集比例: {val_size}")
        
        # 第一次分割：训练集 + 临时集
        X = df['text'].values
        y = df['intent'].values
        
        X_train, X_temp, y_train, y_temp = train_test_split(
            X, y, 
            test_size=test_size,
            random_state=random_state,
            stratify=y
        )
        
        # 第二次分割：验证集 + 测试集  
        X_val, X_test, y_val, y_test = train_test_split(
            X_temp, y_temp,
            test_size=val_size,
            random_state=random_state,
            stratify=y_temp
        )
        
        # 创建数据框
        train_df = pd.DataFrame({'text': X_train, 'intent': y_train})
        val_df = pd.DataFrame({'text': X_val, 'intent': y_val})
        test_df = pd.DataFrame({'text': X_test, 'intent': y_test})
        
        logger.info(f"数据分割完成 - 训练: {len(train_df)}, 验证: {len(val_df)}, 测试: {len(test_df)}")
        
        return train_df, val_df, test_df


def create_sample_data(output_path: str = "data/sample/intent_data.csv") -> pd.DataFrame:
    """
    创建示例数据
    
    Args:
        output_path: 输出文件路径
        
    Returns:
        pd.DataFrame: 示例数据
    """
    sample_data = [
        # 问候类
        ("你好", "问候"),
        ("早上好", "问候"), 
        ("晚上好", "问候"),
        ("您好", "问候"),
        ("hi", "问候"),
        
        # 再见类
        ("再见", "再见"),
        ("拜拜", "再见"),
        ("下次见", "再见"),
        ("goodbye", "再见"),
        ("bye", "再见"),
        
        # 查询天气类
        ("今天天气怎么样", "查询天气"),
        ("明天会下雨吗", "查询天气"),
        ("北京的天气如何", "查询天气"),
        ("查一下天气预报", "查询天气"),
        ("温度多少度", "查询天气"),
        
        # 订票类
        ("我要订票", "订票"),
        ("帮我订一张去北京的机票", "订票"),
        ("订火车票", "订票"),
        ("买票", "订票"),
        ("预订机票", "订票"),
        
        # 播放音乐类
        ("播放音乐", "播放音乐"),
        ("我想听歌", "播放音乐"),
        ("放首歌", "播放音乐"),
        ("播放周杰伦的歌", "播放音乐"),
        ("来首轻音乐", "播放音乐"),
        
        # 设置提醒类
        ("提醒我开会", "设置提醒"),
        ("设置闹钟", "设置提醒"),
        ("明天下午3点提醒我", "设置提醒"),
        ("提醒我买菜", "设置提醒"),
        ("设个提醒", "设置提醒"),
        
        # 搜索类
        ("搜索Python教程", "搜索"),
        ("查找资料", "搜索"),
        ("帮我搜索一下", "搜索"),
        ("百度一下", "搜索"),
        ("找找看", "搜索"),
        
        # 其他类
        ("我想吃饭", "其他"),
        ("无聊", "其他"),
        ("随便聊聊", "其他"),
        ("什么时候", "其他"),
        ("不知道", "其他"),
    ]
    
    # 扩展数据
    extended_data = []
    for text, intent in sample_data:
        extended_data.append((text, intent))
        # 添加一些变体
        if "我" in text:
            extended_data.append((text.replace("我", "我们"), intent))
        if "吗" in text:
            extended_data.append((text.replace("吗", ""), intent))
    
    df = pd.DataFrame(extended_data, columns=['text', 'intent'])
    
    # 保存文件
    output_dir = Path(output_path).parent
    output_dir.mkdir(parents=True, exist_ok=True)
    df.to_csv(output_path, index=False, encoding='utf-8')
    
    logger.info(f"示例数据已保存到: {output_path}")
    return df


if __name__ == "__main__":
    # 创建示例数据用于测试
    sample_df = create_sample_data()
    print(f"创建了 {len(sample_df)} 条示例数据")
    
    # 测试数据加载器
    loader = IntentDataLoader()
    loader.validate_data(sample_df)
    stats = loader.get_data_statistics(sample_df)
    
    # 测试数据分割
    train_df, val_df, test_df = loader.split_data(sample_df)
    print(f"数据分割完成: 训练{len(train_df)}, 验证{len(val_df)}, 测试{len(test_df)}") 